#include "kws.h"
#include "arm_nnfunctions.h"
#include "arm_nnsupportfunctions.h"
#include "local_NN.h"
#include "fsl_debug_console.h"

q7_t* scratch_pad;
q7_t* col_buffer;
q7_t* buffer1;
q7_t* buffer2;

static q7_t *mfcc_buffer;
static q7_t *output;

static MFCC_STR mfcc_str;

const q7_t conv1_wt[CONV1_OUT_CH*CONV1_KX*CONV1_KY]=CONV1_WT;
const q7_t conv1_bias[CONV1_OUT_CH]=CONV1_BIAS;
const q7_t conv2_ds_wt[CONV1_OUT_CH*CONV2_DS_KX*CONV2_DS_KY]=CONV2_DS_WT;
const q7_t conv2_ds_bias[CONV1_OUT_CH]=CONV2_DS_BIAS;
const q7_t conv2_pw_wt[CONV2_OUT_CH*CONV1_OUT_CH]=CONV2_PW_WT;
const q7_t conv2_pw_bias[CONV2_OUT_CH]=CONV2_PW_BIAS;
const q7_t conv3_ds_wt[CONV2_OUT_CH*CONV3_DS_KX*CONV3_DS_KY]=CONV3_DS_WT;
const q7_t conv3_ds_bias[CONV2_OUT_CH]=CONV3_DS_BIAS;
const q7_t conv3_pw_wt[CONV3_OUT_CH*CONV2_OUT_CH]=CONV3_PW_WT;
const q7_t conv3_pw_bias[CONV3_OUT_CH]=CONV3_PW_BIAS;
const q7_t conv4_ds_wt[CONV3_OUT_CH*CONV4_DS_KX*CONV4_DS_KY]=CONV4_DS_WT;
const q7_t conv4_ds_bias[CONV3_OUT_CH]=CONV4_DS_BIAS;
const q7_t conv4_pw_wt[CONV4_OUT_CH*CONV3_OUT_CH]=CONV4_PW_WT;
const q7_t conv4_pw_bias[CONV4_OUT_CH]=CONV4_PW_BIAS;
const q7_t conv5_ds_wt[CONV4_OUT_CH*CONV5_DS_KX*CONV5_DS_KY]=CONV5_DS_WT;
const q7_t conv5_ds_bias[CONV4_OUT_CH]=CONV5_DS_BIAS;
const q7_t conv5_pw_wt[CONV5_OUT_CH*CONV4_OUT_CH]=CONV5_PW_WT;
const q7_t conv5_pw_bias[CONV5_OUT_CH]=CONV5_PW_BIAS;
const q7_t final_fc_wt[CONV5_OUT_CH*OUT_DIM]=FINAL_FC_WT;
const q7_t final_fc_bias[OUT_DIM]=FINAL_FC_BIAS;


int32_t get_top_class(q7_t* prediction)
{
	int32_t max_ind=0;
	int32_t max_val=-128;
	for(int32_t i=0;i<OUT_DIM;i++) {
		if(max_val<prediction[i]) {
			max_val = prediction[i];
			max_ind = i;
		}    
	}
	return max_ind;
}

void ds_cnn_init(void)
{

	mfcc_init(&mfcc_str, NUM_MFCC_COEFFS, FRAME_LEN, MFCC_DEC_BITS);
	mfcc_buffer = (q7_t*)malloc(sizeof(q7_t) * NUM_FRAMES * NUM_MFCC_COEFFS);
	if (mfcc_buffer == NULL)
	{
		PRINTF("The mfcc_buffer malloc failed\r\n");
	}
	output = (q7_t*)malloc(sizeof(q7_t) * OUT_DIM);
	if (output == NULL)
	{
		PRINTF("The output malloc failed\r\n");
	}

	scratch_pad = (q7_t *)malloc(sizeof(q7_t) * SCRATCH_BUFFER_SIZE);
	if (scratch_pad == NULL)
	{
		PRINTF("scratch_pad malloc failed\r\n");
	}
	buffer1 = scratch_pad;
	buffer2 = buffer1 + (CONV1_OUT_CH*CONV1_OUT_X*CONV1_OUT_Y);
	col_buffer = buffer2 + (CONV2_OUT_CH*CONV2_OUT_X*CONV2_OUT_Y);
}

void extract_features(MFCC_STR *mfcc_str, const int16_t *audio_data, q7_t *mfcc_buffer)
{
	uint16_t f;
	int32_t mfcc_buffer_head = 0;
	for (f = 0; f < NUM_FRAMES; f++) 
	{
		mfcc_compute(mfcc_str, audio_data + (f * FRAME_SHIFT),&mfcc_buffer[mfcc_buffer_head]);
		mfcc_buffer_head += NUM_MFCC_COEFFS;
	}
}

void run_nn_to_classify(q7_t* in_data, q7_t* out_data)
{
	//卷积层1
	arm_convolve_HWC_q7_basic_nonsquare(in_data, CONV1_IN_X, CONV1_IN_Y, 1, conv1_wt, CONV1_OUT_CH, CONV1_KX, CONV1_KY, CONV1_PX, CONV1_PY, CONV1_SX, CONV1_SY, conv1_bias, CONV1_BIAS_LSHIFT, CONV1_OUT_RSHIFT, buffer1, CONV1_OUT_X, CONV1_OUT_Y, (q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer1,CONV1_OUT_X*CONV1_OUT_Y*CONV1_OUT_CH);

	//卷积层2
	arm_depthwise_separable_conv_HWC_q7_nonsquare(buffer1,CONV2_IN_X,CONV2_IN_Y,CONV1_OUT_CH,conv2_ds_wt,CONV1_OUT_CH,CONV2_DS_KX,CONV2_DS_KY,CONV2_DS_PX,CONV2_DS_PY,CONV2_DS_SX,CONV2_DS_SY,conv2_ds_bias,CONV2_DS_BIAS_LSHIFT,CONV2_DS_OUT_RSHIFT,buffer2,CONV2_OUT_X,CONV2_OUT_Y,(q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer2,CONV2_OUT_X*CONV2_OUT_Y*CONV2_OUT_CH);

	arm_convolve_1x1_HWC_q7_fast_nonsquare(buffer2, CONV2_OUT_X, CONV2_OUT_Y, CONV1_OUT_CH, conv2_pw_wt, CONV2_OUT_CH, 1, 1, 0, 0, 1, 1, conv2_pw_bias, CONV2_PW_BIAS_LSHIFT, CONV2_PW_OUT_RSHIFT, buffer1, CONV2_OUT_X, CONV2_OUT_Y, (q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer1,CONV2_OUT_X*CONV2_OUT_Y*CONV2_OUT_CH);

	arm_depthwise_separable_conv_HWC_q7_nonsquare(buffer1,CONV3_IN_X,CONV3_IN_Y,CONV2_OUT_CH,conv3_ds_wt,CONV2_OUT_CH,CONV3_DS_KX,CONV3_DS_KY,CONV3_DS_PX,CONV3_DS_PY,CONV3_DS_SX,CONV3_DS_SY,conv3_ds_bias,CONV3_DS_BIAS_LSHIFT,CONV3_DS_OUT_RSHIFT,buffer2,CONV3_OUT_X,CONV3_OUT_Y,(q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer2,CONV3_OUT_X*CONV3_OUT_Y*CONV3_OUT_CH);

	arm_convolve_1x1_HWC_q7_fast_nonsquare(buffer2, CONV3_OUT_X, CONV3_OUT_Y, CONV2_OUT_CH, conv3_pw_wt, CONV3_OUT_CH, 1, 1, 0, 0, 1, 1, conv3_pw_bias, CONV3_PW_BIAS_LSHIFT, CONV3_PW_OUT_RSHIFT, buffer1, CONV3_OUT_X, CONV3_OUT_Y, (q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer1,CONV3_OUT_X*CONV3_OUT_Y*CONV3_OUT_CH);

	arm_depthwise_separable_conv_HWC_q7_nonsquare(buffer1,CONV4_IN_X,CONV4_IN_Y,CONV3_OUT_CH,conv4_ds_wt,CONV3_OUT_CH,CONV4_DS_KX,CONV4_DS_KY,CONV4_DS_PX,CONV4_DS_PY,CONV4_DS_SX,CONV4_DS_SY,conv4_ds_bias,CONV4_DS_BIAS_LSHIFT,CONV4_DS_OUT_RSHIFT,buffer2,CONV4_OUT_X,CONV4_OUT_Y,(q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer2,CONV4_OUT_X*CONV4_OUT_Y*CONV4_OUT_CH);

	arm_convolve_1x1_HWC_q7_fast_nonsquare(buffer2, CONV4_OUT_X, CONV4_OUT_Y, CONV3_OUT_CH, conv4_pw_wt, CONV4_OUT_CH, 1, 1, 0, 0, 1, 1, conv4_pw_bias, CONV4_PW_BIAS_LSHIFT, CONV4_PW_OUT_RSHIFT, buffer1, CONV4_OUT_X, CONV4_OUT_Y, (q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer1,CONV4_OUT_X*CONV4_OUT_Y*CONV4_OUT_CH);

	arm_depthwise_separable_conv_HWC_q7_nonsquare(buffer1,CONV5_IN_X,CONV5_IN_Y,CONV4_OUT_CH,conv5_ds_wt,CONV4_OUT_CH,CONV5_DS_KX,CONV5_DS_KY,CONV5_DS_PX,CONV5_DS_PY,CONV5_DS_SX,CONV5_DS_SY,conv5_ds_bias,CONV5_DS_BIAS_LSHIFT,CONV5_DS_OUT_RSHIFT,buffer2,CONV5_OUT_X,CONV5_OUT_Y,(q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer2,CONV5_OUT_X*CONV5_OUT_Y*CONV5_OUT_CH);

	arm_convolve_1x1_HWC_q7_fast_nonsquare(buffer2, CONV5_OUT_X, CONV5_OUT_Y, CONV4_OUT_CH, conv5_pw_wt, CONV5_OUT_CH, 1, 1, 0, 0, 1, 1, conv5_pw_bias, CONV5_PW_BIAS_LSHIFT, CONV5_PW_OUT_RSHIFT, buffer1, CONV5_OUT_X, CONV5_OUT_Y, (q15_t*)col_buffer, NULL);
	arm_relu_q7(buffer1,CONV5_OUT_X*CONV5_OUT_Y*CONV5_OUT_CH);

	arm_avepool_q7_HWC_nonsquare (buffer1,CONV5_OUT_X,CONV5_OUT_Y,CONV5_OUT_CH,CONV5_OUT_X,CONV5_OUT_Y,0,0,1,1,1,1,NULL,buffer2, 2);

	arm_fully_connected_q7(buffer2, final_fc_wt, CONV5_OUT_CH, OUT_DIM, FINAL_FC_BIAS_LSHIFT, FINAL_FC_OUT_RSHIFT, final_fc_bias, out_data, (q15_t*)col_buffer);
}

int32_t kws_using_ds_cnn(int16_t *audio_buf)
{
	extract_features(&mfcc_str, audio_buf, mfcc_buffer);
	run_nn_to_classify(mfcc_buffer, output);
	arm_softmax_q7(output, OUT_DIM, output);
	return get_top_class(output);
}

